/*
 *
 * Copyright (c) 2006, 2007, 2008 Prabhu Gurumurthy <pgurumur@gmail.com>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include "trie.h"

static char _PyTrie_ConnectDoc[] = "Red Black tree implementation using C";

PyTypeObject PyTrieObjectType = {
   PyObject_HEAD_INIT(NULL)
      0,
   "trie",
   sizeof(PyTrieObject),
   0,

   (destructor)_PyTrie_ObjectDeAlloc,
   NULL,
   (getattrfunc)_PyTrie_GetAttr,
   (setattrfunc)_PyTrie_SetAttr,
   NULL,
   NULL,

   NULL,
   NULL,
   NULL,

   NULL,
   NULL,
   NULL,
   NULL,
   NULL,

   NULL,

   Py_TPFLAGS_DEFAULT,
   (char *)_PyTrie_ConnectDoc,
   (traverseproc)_PyTrie_Traverse,

   NULL,

   NULL,

   0,

   NULL,
   NULL,
   (struct PyMethodDef *)PyTrie_Methods,
   (TQ_MEMBER(*))_PyTrieObjectMemberList,
   NULL,
   NULL,

   NULL,
   0,
   NULL,
   0,

   NULL,
   NULL,
   NULL,
   NULL,
   NULL,
   NULL,
   NULL,
   NULL,
   0,
   0,
   0
};

DL_EXPORT(void)
inittrie(void)
{
   PyObject *module, *dict, *version;

   module = Py_InitModule4("trie", PyTrie_Methods, _PyTrie_Doc,
         (PyObject *)NULL, PYTHON_API_VERSION);

   if ((dict = PyModule_GetDict(module)) == NULL)
      Py_FatalError("PyModule_GetDict()");

   version = PyString_FromString(_PyTrie_Version);
   if (version == NULL || PyDict_SetItemString(dict, _PyTrie_Version, version))
      Py_FatalError("could not initialize version string");

   Py_DECREF(version);

   PyTrie_Error = Py_BuildValue("s", "TrieError");
   if (PyTrie_Error == NULL || 
         PyDict_SetItemString(dict, "TrieError", PyTrie_Error))
      Py_FatalError("could not initialize trie");

   Py_DECREF(PyTrie_Error);

   PyTrieObjectType.ob_type = &PyType_Type;
   PyTrieObjectType.tp_alloc = PyType_GenericAlloc;
   PyTrieObjectType.tp_new = PyType_GenericNew;
   PyTrieObjectType.tp_free = _PyObject_GC_Del;

   if (PyDict_SetItemString(dict, "TrieObject", (PyObject *)&PyTrieObjectType))
      Py_FatalError("could not initialize trie object type");

   Py_INCREF(&PyTrieObjectType);
   fflush(stdout);
}

static void
_PyTrie_ObjectDeAlloc(PyTrieObject *self)
{
   if (self->closed == 0) {
      PyObject *object = PyTrie_Empty(self, NULL);
      Py_XDECREF(object);
   }
}

static int
_PyTrie_Traverse(PyTrieObject *self, visitproc visit, void *arg)
{
   int retval = 0;

   if (self->converter) {
      int ix = 0;
      if (!(ix = visit(self->converter, arg)))
         retval = ix;
   }

   return retval;
}

static PyObject *
_PyTrie_GetAttr(PyTrieObject *self, char *name)
{
   PyObject *retval = Py_FindMethod(PyTrie_Methods, (PyObject *)self, name);
   if (retval == NULL)
      PyErr_SetString(PyExc_AttributeError, name);

   return retval;
}

static int
_PyTrie_SetAttr(PyTrieObject *self, char *name, PyObject *object)
{
   if (object == NULL) {
      PyErr_SetString(PyExc_AttributeError, "cannot delete attributes");
      return -1;
   }

   TQ_MEMBER(*list);
   for (list = _PyTrieObjectMemberList; list->name; list++)
      if (strncmp(list->name, name, strlen(name)) == 0)
         return PyMember_SetOne((char *)self, list, object);

   PyErr_SetString(PyExc_AttributeError, name);
   return -1;
}

static void
PyTrie_SetError(int isexit, const char *fmt, ...)
{
   va_list ap;
   char errmesg[BUFSIZ];

   memset(errmesg, '\0', BUFSIZ);
   va_start(ap, fmt);
   vsnprintf(errmesg, BUFSIZ, fmt, ap);
   va_end(ap);

   PyErr_SetString(PyTrie_Error, errmesg);
   if (isexit)
      Py_Exit(1);
}

static int
PyTrie_Compare(struct trie *first, struct trie *second)
{
   int retval;
   int cmp = PyObject_Cmp(first->value, second->value, &retval);
   cmp = cmp;
   return retval;
}

static PyObject *
PyTrie_Init(PyObject *self, PyObject *args)
{
   PyTrieObject *object = TQ_ALLOC(PyTrieObject, PyTrieObjectType);
   if (object == NULL) 
      PyTrie_SetError(1, "cannot initialize trie object");

   self = self;
   object->count = 0;
   object->closed = 0;
   RB_INIT(&object->head);
   Py_INCREF(object);

   return (PyObject *)object;
}

RB_PROTOTYPE(treehead, trie, entry, PyTrie_Compare);
RB_GENERATE(treehead, trie, entry, PyTrie_Compare);

static PyObject *
PyTrie_Empty(PyTrieObject *self, PyObject *args)
{
   PyObject *retval;
   if (self) {
      if ((self->closed == 0) && (self->count >= 1)) {
         PyObject *input;
         if (args) {
            int ok = PyArg_ParseTuple(args, "O", &input);
            if (ok) {
               struct trie find;
               find.value = input;
               struct trie *node = RB_FIND(treehead, &self->head, &find);
               if (node) {
                  retval = node->value;
                  RB_REMOVE(treehead, &self->head, node);
                  free(node);
                  self->count--;
               }
            }
         }
         else {
            struct trie *node;
            RB_FOREACH(node, treehead, &(self->head)) {
               RB_REMOVE(treehead, &self->head, node);
               free(node);
               self->count--;
            }

            self->closed = 1;

            Py_DECREF(self);
            Py_INCREF(Py_None);
            retval = Py_None;
         }
      }
   }

   return retval;
}

static PyObject *
PyTrie_Add(PyTrieObject *self, PyObject *args)
{
   if (self) {
      PyObject *input;
      if (args) {
         int ok = PyArg_ParseTuple(args, "O", &input);
         if (ok) {
            struct trie *node = malloc(sizeof(struct trie));
            if (node) {
               node->value = input;
               RB_INSERT(treehead, &self->head, node);
               self->count++;
            }
            else
               PyTrie_SetError(1, strerror(errno));
         }
      }
      else
         PyTrie_SetError(0, "no items provided to add");
   }
   else
      PyTrie_SetError(0, "initialize trie first");

   Py_INCREF(Py_None);
   return Py_None;
}

static PyObject *
PyTrie_Find(PyTrieObject *self, PyObject *args)
{
   PyObject *retval;

   if (self && (self->count >= 1) && (self->closed == 0)) {
      if (args) {
         int ok = PyArg_ParseTuple(args, "O", &retval);
         if (ok) {
            struct trie find;
            find.value = retval;
            struct trie *node = RB_FIND(treehead, &self->head, &find);
            if (node)
               retval = node->value;
         }
         else {
            Py_INCREF(Py_None);
            retval = Py_None;
         }
      }
      else {
         Py_INCREF(Py_None);
         retval = Py_None;
      }
   }
   else
      PyTrie_SetError(0, "initialize trie first!");

   return retval;
}

static PyObject *
PyTrie_Remove(PyTrieObject *self, PyObject *args)
{
   return PyTrie_Empty(self, args);
}

static PyObject *
PyTrie_List(PyTrieObject *self, PyObject *args)
{
   PyObject *retval = PyList_New(0);
   if (self && (self->count >= 1) && (self->closed == 0)) {
      struct trie *node;
      RB_FOREACH(node, treehead, &self->head) {
         if (PyList_Append(retval, node->value) == -1)
            PyTrie_SetError(0, "unable to insert item to a list");
      }
   }

   return retval;
}

static PyObject *
PyTrie_Count(PyTrieObject *self, PyObject *args)
{
   PyObject *retval = NULL;
   int temp = 0;
   if (self && self->closed == 0)
      retval = Py_BuildValue("i", self->count);
   else
      retval = Py_BuildValue("i", temp);

   return retval;
}
